# 导入函数库
from jqdata import *
import math

'''
策略思路：
1.筛选出当天涨停数最多的行业板块，以及该板块的涨停股作为第二天备选
2.备选股中，开盘没涨停的，在盘中涨停的时候买入（主要开盘涨停的一般买不到，盘中利用同花顺的条件单试试）
3.持仓股尾盘没涨停就出了

关于回测：
没有盘中实时监控委托，先尾盘强制成交看看效果，误差可能较大（1、真实不一定能买进 2、这样过滤了涨停开板回落的票，使得收益比实际更高）
'''


def initialize(context):
    set_benchmark('000300.XSHG')
    set_option('use_real_price', True)
    # 过滤掉order系列API产生的比error级别低的log
    log.set_level('order', 'error')
    # 避免未来数据
    set_option("avoid_future_data", True)

    # 强制成交，
    # 没有盘中实时监控委托，先尾盘强制成交看看效果，误差可能较大（1、真实不一定能买进 2、这样过滤了涨停开板回落的票，使得收益比实际更高）
    set_option("match_by_signal", True)

    set_order_cost(OrderCost(close_tax=0.001, open_commission=0.00015, close_commission=0.00015, min_commission=5),
                   type='stock')

    # 备选股
    g.option_stocks = []

    run_daily(daily_trade, time='14:50', reference_security='000300.XSHG')
    # 2022年至今回测结果：按概念比按行业收益高不少
    # run_daily(get_industry_option_stocks, time='17:00', reference_security='000300.XSHG')
    run_daily(get_concept_option_stocks, time='17:00', reference_security='000300.XSHG')


def daily_trade(context):
    # 先卖后买
    tail_sell(context)
    buy_trade(context)


# 买操作
def buy_trade(context):
    # 每次买入单票最多3成仓
    buy_max_per_stock = round(context.portfolio.total_value * 0.3, 2)
    current_data = get_current_data()
    for stock in g.option_stocks:
        available_cash = context.portfolio.available_cash
        if available_cash < 5000:
            return
        buy_cash = min(available_cash, buy_max_per_stock)
        current_stock = current_data[stock]
        current_price = current_stock.last_price
        day_open = current_stock.day_open
        high_limit = current_stock.high_limit
        # 开盘没涨停，但现在涨停了，符合条件，买
        if day_open < high_limit and current_price == high_limit:
            order_value(stock, buy_cash)


# 尾盘卖出
def tail_sell(context):
    # 持仓没涨停，就卖了
    long_positions_dict = context.portfolio.long_positions
    current_data = get_current_data()
    for position in list(long_positions_dict.values()):
        current_stock = current_data[position.security]
        current_price = current_stock.last_price
        day_open = current_stock.day_open
        high_limit = current_stock.high_limit
        if current_price < high_limit:
            order_target(position.security, 0)


# ================================按行业选股 start==========================

# 获取行业备选股
def get_industry_option_stocks(context):
    # 今天涨停的票
    today_limit_stocks = get_today_limit_stocks(context)
    # 每个涨停票所属行业
    stock_industry_df = get_stock_industry_df(context, today_limit_stocks)

    # 每个行业涨停股票数
    # 加“count”字段只是为了更清晰，其实没必要加这个count, 因为groupby+count后，除分组字段，其他字段值都是count数量
    stock_industry_df["count"] = 1
    # 按申万二级行业试试
    industry_count_df = stock_industry_df.groupby(["sw_L2"]).count()
    industry_count_df = industry_count_df.sort_values(by="count", ascending=False)
    # 涨停数最多的行业
    max_industry = industry_count_df.index[0]

    option_stocks_df = stock_industry_df[stock_industry_df["sw_L2"] == max_industry]
    options_stocks = list(option_stocks_df["code"])

    log.info("==============涨停数最多的申万行业：%s, 具体涨停股：%s" % (max_industry, str(options_stocks)))
    g.option_stocks = options_stocks


# 获取股票所属行业
def get_stock_industry_df(context, stocks):
    stock_industry_dict = get_industry(stocks, date=current_dt)
    # 申万一、二、三级行业
    sw_L1 = []
    sw_L2 = []
    sw_L3 = []

    for stock in stocks:
        industry_dict = stock_industry_dict[stock]
        '''
        {'jq_l1': {'industry_code': 'HY007', 'industry_name': '金融'},
         'sw_l2': {'industry_code': '801783', 'industry_name': '股份制银行II'},
         'sw_l1': {'industry_code': '801780', 'industry_name': '银行I'},
         'jq_l2': {'industry_code': 'HY07101', 'industry_name': '综合性银行'},
         'sw_l3': {'industry_code': '857831', 'industry_name': '股份制银行III'},
         'zjw': {'industry_code': 'J66', 'industry_name': '货币金融服务'}}
        '''
        sw_L1.append(industry_dict["sw_l1"]["industry_name"])
        sw_L2.append(industry_dict["sw_l2"]["industry_name"])
        sw_L3.append(industry_dict["sw_l3"]["industry_name"])

    stock_industry_df = pd.DataFrame(columns=['code', 'sw_L1', 'sw_L2', 'sw_L3'])
    stock_industry_df["code"] = stocks
    stock_industry_df["sw_L1"] = sw_L1
    stock_industry_df["sw_L2"] = sw_L2
    stock_industry_df["sw_L3"] = sw_L3

    return stock_industry_df


# ================================按行业选股 end==========================


# ================================按概念选股 start==========================

def get_concept_option_stocks(context):
    # 今天涨停的票
    today_limit_stocks = get_today_limit_stocks(context)
    # 每个涨停票所属概念，一个股票属于多个概念，则每个概念一行
    stock_concept_df = get_stock_concept_df(context, today_limit_stocks)

    # 每个概念涨停股票数
    stock_concept_df["count"] = 1
    concept_count_df = stock_concept_df.groupby(["concept"]).count()
    concept_count_df = concept_count_df.sort_values(by="count", ascending=False)

    max_concept = concept_count_df.index[0]
    option_stocks_df = stock_concept_df[stock_concept_df["concept"] == max_concept]
    # 因为每个股票的每个概念都是独立一行，所以需要去重
    options_stocks = set(option_stocks_df["code"])

    print("涨停数最多的概念：%s, 具体涨停股：%s" % (max_concept, str(options_stocks)))
    g.option_stocks = options_stocks
    return options_stocks


# 获取股票所属概念（多个概念）
def get_stock_concept_df(context, stocks):
    '''
    {'002903.XSHE': {'jq_concept': [
        {'concept_code': 'SC0005','concept_name': '华为'},
        {'concept_code': 'SC0019', 'concept_name': '苹果'},
        {'concept_code': 'SC0030', 'concept_name': '小米'},
        {'concept_code': 'SC0031', 'concept_name': '富士康'},
        {'concept_code': 'SC0053', 'concept_name': '机器人'},
        {'concept_code': 'SC0099', 'concept_name': '工业4.0'},
        {'concept_code': 'SC0224', 'concept_name': '专精特新'},
        {'concept_code': 'SC0229', 'concept_name': '工业母机'},
        {'concept_code': 'SC0284', 'concept_name': '比亚迪概念'}
    ]}}
    '''
    stock_concept_dict = get_concept(stocks, date=context.current_dt)
    # 股票所属概念 [["002240.XSHE", "华为"], ["002240.XSHE", "苹果"]]
    stock_concept = []
    # 排除一些共性的概念
    exclude_concept = {"沪股通", "深股通", "转融券标的", "融资融券", "华为", "苹果", "特斯拉"}

    for stock in stocks:
        concept_list = stock_concept_dict[stock]["jq_concept"]
        for concept in concept_list:
            if concept["concept_name"] not in exclude_concept:
                stock_concept.append([stock, concept["concept_name"]])

    # 转换为dataframe
    stock_concept_df = pd.DataFrame(stock_concept, columns=['code', 'concept'])

    return stock_concept_df


# ================================按概念选股 end==========================


# 获取当天涨停的票，过滤掉上市不足半年的票
def get_today_limit_stocks(context):
    # 获取所有股票，排除上市不足半年
    all_stocks = get_all_securities(types=['stock'], date=context.current_dt)
    # 半年前的日期
    pre_half_year_date = context.current_dt.date() - timedelta(days=180)
    # 过滤上市不足半年的股票
    all_stocks = all_stocks[all_stocks['start_date'] < pre_half_year_date]

    today_df = get_price(list(all_stocks.index), end_date=context.current_dt, count=1, frequency='1d',
                         fields=['close', 'high_limit'], panel=False, fill_paused=False)
    today_limit_df = today_df[today_df['close'] == today_df['high_limit']]

    return list(today_limit_df.code)










